#include <cmath>
#include <assert.h>
#include <algorithm>
#include <complex>
#include <vector>

template <class T>
void FFT(std::vector<std::complex<T> >& data,bool inverse=false)
{
    // Converts data to it's Fourier transform
    // data must be an integral power of two in length
    const double PI = 3.1415926535897973;

    const uint n=data.size();

    // First bit reverse the data vector
    // eg. for n = 16:
    // Swap:  0001 <-> 1000  (1,8)
    //        0010 <-> 0100  (2,4)
    //        0011 <-> 1100  (3,12)
    //        0101 <-> 1010  (5,10)
    //        0111 <-> 1110  (7,14)
    for (uint i=0, j=0;i<n/2;i++) {
        // j is bit-reversed version of i
        assert(j<n);
        if (j>i) std::swap(data[i],data[j]);
        // Find new j for i+1
        uint m = n/2;
        while (j>=m) {
            assert(m>1); 
            j -= m;
            m /= 2;
        }
        j += m;
    }

    // Now do the Danielson-Lanczos recursion:
    for(uint m = 1; m < n; m *= 2) {
        // At the start of the loop, all Fourier transforms
        // in blocks of size m are done.
        // For the first pass, each 1-element Fourier transform 
        // is done due to the bit reversal above
        // When m == n, don't do the loop, since whole 
        // vector is done.
        const std::complex<T> W = inverse ?
            std::polar(1.,-PI/m) : 
            std::polar(1.,PI/m);
        std::complex<T> Wk = 1.; // Wk = W^k
        for (uint k=0;k<m;k++,Wk*=W) {
            for (uint i=k;i<n;i+=2*m) {
                // D-J Lemma:
                // Fk = Fk^even + W^k Fk^odd
                // Do this for each adjacent pair of blocks
                // to get new blocks 2x the size.
                // This loop is over each pair of blocks.
                // Outer loop is over location in the block.
                //
                // Note: Fk^e and Fk^o are periodic with period m
                // and W^(k+m) = -W^k
                // So: F_k+m = Fk^e - W^k Fk^o
                //
                uint j=i+m;
                std::complex<T> temp = Wk*data[j];
                data[j] = data[i] - temp;
                data[i] += temp;
            }
        }
    }
    if (inverse) {
        for(uint i=0;i<n;i++) data[i]/=n;
    }
}

template <class T>
void FFT(
    const std::vector<T>& realdata, std::vector<std::complex<T> >& transform)
{
    // Takes the Fourier transform of a real vector
    // It returns only the positive half in transform, since 
    // it is symmetric about 0.
    // NOTE: realdata must have a size n equal to an integral power of 2.
    // Transform will be returned with size n/2 + 1.
    // Basically, it is just a tricky way to pack the data into half the space
    // for a call to the usual fft.
    const double PI = 3.1415926535897973;
    const uint n = realdata.size()/2;

    // Split real vector f into even and odd parts, fe and fo
    // Take the fft of fe + i fo
    transform.reserve(n+1);
    transform.resize(n);
    for(uint i=0;i<n;i++) transform[i] = 
        std::complex<T>(realdata[2*i],realdata[2*i+1]);
    FFT(transform);

    // The FT of fe has the symmetry F_n-k = F_k^*
    // The FT of (i fo) has the symmetry (i G_n-k) = -(i G_k)^*
    // Let's call the FT of this combinations H_k = F_k + (i G_k)
    // Then F_k = (H_k + H_n-k^*)/2
    // And G_k = (H_k - H_n-k^*)/2i
    // The Danielson-Lanczos lemma says that our answer T is
    // T_k = F_k + W^k G_k
    const std::complex<T> W = std::polar(1.,PI/n);
    std::complex<T> Wkinv2i = std::complex<T>(0.,-0.5)*W; // = W^k/2i
    for (uint k=1;k<n/2;k++,Wkinv2i*=W) {
        // Deal with k=0,n/2,n below
        std::complex<T> Fk = (transform[k]+conj(transform[n-k]))*0.5;
        std::complex<T> WkGk = (transform[k]-conj(transform[n-k]))*Wkinv2i;
        transform[k] = Fk + WkGk;
        // Note: W^(n-k) = -W^-k = -conj(W^k)
        transform[n-k] = conj(Fk - WkGk);
    }
    // For k=n/2, W^k = i, F_k = Re(H_k), G_k = Im(H_k)
    // So T_n/2 = H_n/2, so don't need to do anything
    // For k=0,n, need fact that H_n = H_0 from periodicity
    // For k=n, W^k = -1, F_k = Re(H_0), G_k = Im(H_0)
    // So T_0 = Re(H_0) - Im(H_0)
    // For k=0, W^k = 1, F_k = Re(H_0), G_k = Im(H_0)
    // So T_0 = Re(H_0) + Im(H_0)
    T re0 = real(transform[0]);
    T im0 = imag(transform[0]);
    transform.push_back(re0-im0);
    transform[0] = re0+im0;
}

template <class T>
void InvFFT(
    const std::vector<std::complex<T> >& transform, std::vector<T>& realdata)
{
    // Returns the inverse Fourier transform of a transform given the 
    // positive frequency data only, returning the (real) result as realdata.
    // NOTE: transform must be 1 more than an integral power of 2,
    // realdata will be returned with a size of twice this power of 2.
    // Also, transform[0] and transform[n] are assumed to be real
    // (or more accurately, their imaginary parts are ignored).
    const double PI = 3.1415926535897973;
    const uint n = transform.size()-1;
    std::vector<std::complex<T> > work(n);

    // For inverse transform, we need to construct H_k = F_k + iG_k
    // from T_k = F_k + W^k G_k  
    // From the relations for F_k and G_k above:
    // F_k = (T_k + T_N-k*)/2
    // G_k = (T_k - T_N-k*)/2W^k
    const std::complex<T> invW = std::polar(1.,-PI/n);
    std::complex<T> iinv2Wk = std::complex<T>(0,0.5)*invW; // = i/2W^k
    for (uint k=1;k<n/2;k++,iinv2Wk*=invW) {
        // Deal with k=0,n/2 below
        std::complex<T> Fk = (transform[k]+conj(transform[n-k]))*0.5;
        std::complex<T> iGk = (transform[k]-conj(transform[n-k]))*iinv2Wk;
        work[k] = Fk + iGk;
        work[n-k] = conj(Fk - iGk);
    }
    work[n/2] = transform[n/2];
    T F0 = (real(transform[0]) + real(transform[n]))*0.5;
    T G0 = (real(transform[0]) - real(transform[n]))*0.5;
    work[0] = std::complex<T>(F0,G0);

    FFT(work,true);
    for(uint i=0;i<n;i++) {
        realdata[2*i] = real(work[i]);
        realdata[2*i+1] = imag(work[i]);
    }
}


